package com.xzy.mybatis.cipher.plugin;

import cn.hutool.core.collection.CollectionUtil;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.xzy.mybatis.cipher.spring.MybatisCipherContext;
import com.xzy.mybatis.cipher.support.DatasourceHolder;
import com.xzy.mybatis.cipher.support.strategy.parse.ParseStrategy;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.util.ArrayList;
import java.util.List;

/**
 * Mybatis 插件：加密
 *
 * @author xzy.xiao
 * @since 2025/1/14  15:59
 */
@Intercepts({
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})
})
@Slf4j
public class EncryptPlugin implements Interceptor {

    private final DatasourceHolder datasourceHolder;
    private final MybatisCipherContext mybatisCipherContext;

    public EncryptPlugin(DatasourceHolder datasourceHolder, MybatisCipherContext mybatisCipherContext) {
        this.datasourceHolder = datasourceHolder;
        this.mybatisCipherContext = mybatisCipherContext;
    }

    /**
     * Intercepts the execution of Mybatis SQL statements to perform data encryption.
     *
     * @param invocation Invocation information, containing information such as the method to be executed and parameters
     * @return The result of the original method execution
     * @throws Throwable Throws an exception if the execution fails
     */
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        String datasourceName = datasourceHolder.getCurrentDatasourceName();

        // Encrypt and replace SQL parameters
        encrypt(datasourceName, invocation);

        // Continue the execution of the original method
        return invocation.proceed();
    }

    // SUPPORT

    /**
     * Performs encryption processing on SQL parameters.
     *
     * @param invocation Invocation information, containing information such as the method to be executed and parameters
     */
    private void encrypt(String datasourceName, Invocation invocation) {
        Object[] args = invocation.getArgs();

        MappedStatement mappedStatement = (MappedStatement) args[0];
        Configuration configuration = mappedStatement.getConfiguration();

        Object parameter = args[1];
        BoundSql boundSql = mappedStatement.getBoundSql(parameter);

        encrypt(datasourceName, configuration, boundSql);
    }

    /**
     * Encrypts SQL parameters based on the configuration and BoundSql.
     *
     * @param configuration Mybatis configuration object
     * @param boundSql      SQL statement and parameter binding object
     */
    private void encrypt(String datasourceName, Configuration configuration, BoundSql boundSql) {
        Object parameterObject = boundSql.getParameterObject();
        if (null == parameterObject) {
            return;
        }

        List<ParameterMapping> parameterMappingList = boundSql.getParameterMappings();
        if (CollectionUtil.isEmpty(parameterMappingList)) {
            return;
        }

        // 1.Extract property value
        MetaObject metaObject = configuration.newMetaObject(parameterObject);
        List<Object> originalPropertyValueList = new ArrayList<>();
        for (ParameterMapping parameterMapping : parameterMappingList) {
            String propertyName = parameterMapping.getProperty();
            if (metaObject.hasGetter(propertyName)) {
                originalPropertyValueList.add(metaObject.getValue(propertyName));
            } else if (boundSql.hasAdditionalParameter(propertyName)) {
                originalPropertyValueList.add(boundSql.getAdditionalParameter(propertyName));
            } else {
                originalPropertyValueList.add(metaObject.getOriginalObject());
            }
        }

        // 2.Encrypt property value
        String sql = boundSql.getSql();
        SQLStatement sqlStatement = SQLUtils.parseStatements(sql, "mysql").get(0);
        ParseStrategy parseStrategy = mybatisCipherContext.getSqlParseStrategyRegistry().getStrategy(sqlStatement);
        if (null == parseStrategy) {
            log.warn("Unsupported sql type! sql: {}", sql);
            return;
        }

        List<Object> encryptedPropertyValueList = parseStrategy.parse(datasourceName, sqlStatement, originalPropertyValueList);

        // 3.Replace property value
        for (int i = 0; i < parameterMappingList.size(); i++) {
            String propertyName = parameterMappingList.get(i).getProperty();
            Object encryptedPropertyValue = encryptedPropertyValueList.get(i);

            if (metaObject.hasSetter(propertyName)) {
                metaObject.setValue(propertyName, encryptedPropertyValue);
            } else if (boundSql.hasAdditionalParameter(propertyName)) {
                boundSql.setAdditionalParameter(propertyName, encryptedPropertyValue);
            }
        }

    }

}
